import argparse
import json
import os
import time

import uvicorn
from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile
from fastapi.responses import JSONResponse, StreamingResponse

app = FastAPI()


UPLOAD_DIR = "./uploads"


@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    return JSONResponse(
        status_code=500,
        content={"status": "error", "code": 500, "message": "Internal Server Error", "details": str(exc)},
    )


@app.get("/check")
def check():
    """
    检查是否能正确响应
    """
    try:
        return {"message": "success"}
    except Exception as e:
        return {"message": "error", "details": str(e)}


@app.post("/upload-with-progress")
async def upload_with_progress(file: UploadFile = File(...)):
    """文件上传接口，实时返回上传进度"""
    total_size = 0
    file_path = os.path.join(UPLOAD_DIR, file.filename)

    try:
        with open(file_path, "wb") as buffer:
            while True:
                chunk = await file.read(1024 * 1024)  # 每次读取1MB
                if not chunk:
                    break
                total_size += len(chunk)
                buffer.write(chunk)

                # 计算并返回进度
                progress = min(100, int((total_size / file.size) * 100))
                yield {"status": "uploading", "progress": progress, "received": total_size, "total": file.size}

        yield {"status": "completed", "filename": file.filename, "size": total_size, "saved_path": file_path}
    except Exception as e:
        yield {"status": "error", "message": str(e)}


# 文件上传
@app.post("/upload")
async def upload_with_progress(
    file: UploadFile = File(..., max_size=10 * 1024 * 1024),
    source: str = Form(...),
    target: str = Form(...),
):
    print(file.filename, source, target)
    if file.file.seek(0, 2) > 10 * 1024 * 1024:
        raise HTTPException(status_code=413, detail="File too large, Please upload files no larger than 10Mb.")

    # 重置文件指针
    file.file.seek(0)

    file_path = os.path.join(UPLOAD_DIR, file.filename)

    file_context = await file.read()
    with open(file_path, "wb") as buffer:
        buffer.write(file_context)

    return {"filename": file.filename, "content_type": file.content_type}


if __name__ == "__main__":
    # 启动服务（指定端口和主机，允许外部访问）
    # parse command line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=8000, help="端口号，默认8000")
    args = parser.parse_args()

    # initial
    os.makedirs(UPLOAD_DIR, exist_ok=True)

    # start the FastAPI app
    uvicorn.run("upload_server:app", host="0.0.0.0", port=args.port, reload=True)
    print(f"Server is running on http://localhost:{args.port}")
